#!/usr/bin/env python3
"""
MyLangExtract 简单使用示例
演示如何使用工具进行结构化信息提取
"""
import os
import requests
import json

def extract_person_info():
    """提取人物信息示例"""
    print("🧪 人物信息提取示例")
    print("-" * 30)
    
    # 检查 API 密钥
    api_key = os.environ.get('ZHIPU_API_KEY')
    if not api_key:
        print("❌ 请先设置 ZHIPU_API_KEY 环境变量")
        return
    
    # API 配置
    url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    
    # 定义提取字段
    schema = {
        "name": "extract_person",
        "description": "提取人物信息",
        "parameters": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "姓名"},
                "job": {"type": "string", "description": "职业"},
                "age": {"type": "string", "description": "年龄"},
                "location": {"type": "string", "description": "地点"}
            },
            "required": ["name"]
        }
    }
    
    # 测试文本
    texts = [
        "张三是一名软件工程师，今年28岁，在北京工作",
        "李四医生35岁，在上海的医院工作",
        "王五是老师，30岁，住在广州"
    ]
    
    for i, text in enumerate(texts, 1):
        print(f"\n📝 文本 {i}: {text}")
        
        # 构建请求
        data = {
            'model': 'glm-4',
            'messages': [
                {'role': 'user', 'content': f'请从以下文本中提取人物信息：{text}'}
            ],
            'tools': [{'type': 'function', 'function': schema}],
            'tool_choice': {'type': 'function', 'function': {'name': 'extract_person'}},
            'temperature': 0.1
        }
        
        try:
            # 发送请求
            response = requests.post(url, headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                
                # 解析结果
                if ('choices' in result and len(result['choices']) > 0 and
                    'message' in result['choices'][0] and 
                    'tool_calls' in result['choices'][0]['message']):
                    
                    tool_calls = result['choices'][0]['message']['tool_calls']
                    if tool_calls and len(tool_calls) > 0:
                        arguments = json.loads(tool_calls[0]['function']['arguments'])
                        
                        print("✅ 提取结果:")
                        for key, value in arguments.items():
                            if value:
                                print(f"   {key}: {value}")
                    else:
                        print("❌ 未找到提取结果")
                else:
                    print("❌ 响应格式错误")
            else:
                print(f"❌ API 错误: {response.status_code}")
                
        except Exception as e:
            print(f"❌ 请求异常: {e}")

def extract_company_info():
    """提取公司信息示例"""
    print("\n🏢 公司信息提取示例")
    print("-" * 30)
    
    api_key = os.environ.get('ZHIPU_API_KEY')
    if not api_key:
        print("❌ 请先设置 ZHIPU_API_KEY 环境变量")
        return
    
    url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    
    # 公司信息提取 schema
    schema = {
        "name": "extract_company",
        "description": "提取公司信息",
        "parameters": {
            "type": "object",
            "properties": {
                "company": {"type": "string", "description": "公司名称"},
                "industry": {"type": "string", "description": "行业"},
                "founded": {"type": "string", "description": "成立时间"},
                "location": {"type": "string", "description": "总部地点"},
                "ceo": {"type": "string", "description": "CEO"}
            },
            "required": ["company"]
        }
    }
    
    text = "腾讯控股有限公司是一家中国的互联网公司，成立于1998年，总部位于深圳，现任CEO是马化腾"
    
    print(f"📝 文本: {text}")
    
    data = {
        'model': 'glm-4',
        'messages': [
            {'role': 'user', 'content': f'请从以下文本中提取公司信息：{text}'}
        ],
        'tools': [{'type': 'function', 'function': schema}],
        'tool_choice': {'type': 'function', 'function': {'name': 'extract_company'}},
        'temperature': 0.1
    }
    
    try:
        response = requests.post(url, headers=headers, json=data, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            
            if ('choices' in result and len(result['choices']) > 0 and
                'message' in result['choices'][0] and 
                'tool_calls' in result['choices'][0]['message']):
                
                tool_calls = result['choices'][0]['message']['tool_calls']
                if tool_calls and len(tool_calls) > 0:
                    arguments = json.loads(tool_calls[0]['function']['arguments'])
                    
                    print("✅ 提取结果:")
                    for key, value in arguments.items():
                        if value:
                            print(f"   {key}: {value}")
                else:
                    print("❌ 未找到提取结果")
            else:
                print("❌ 响应格式错误")
        else:
            print(f"❌ API 错误: {response.status_code}")
            
    except Exception as e:
        print(f"❌ 请求异常: {e}")

def main():
    """主函数"""
    print("MyLangExtract 使用示例")
    print("=" * 50)
    
    # 检查环境
    if not os.environ.get('ZHIPU_API_KEY'):
        print("⚠️  请先配置 API 密钥:")
        print("export ZHIPU_API_KEY='your_api_key_here'")
        print("\n然后重新运行此脚本")
        return
    
    # 运行示例
    extract_person_info()
    extract_company_info()
    
    print("\n" + "=" * 50)
    print("🎉 示例运行完成！")
    print("\n💡 提示:")
    print("- 修改 schema 定义不同的提取字段")
    print("- 调整提示词以适应不同的文本类型")
    print("- 查看 QUICKSTART.md 了解更多用法")

if __name__ == "__main__":
    main()
